import json
import os
import sys
from story_generator import StoryGenerator
from image_prompt_generator import ImagePromptGenerator
from image_generator import ImageGenerator
from utils.file_utils import save_to_file

class StorybookApp:
    def __init__(self):
        self.config = self.load_config()
        self.story_generator = StoryGenerator(self.config)
        self.image_prompt_generator = ImagePromptGenerator(self.config)
        self.image_generator = ImageGenerator(self.config)
        
    def load_config(self):
        """加载配置文件"""
        try:
            config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'config.json')
            with open(config_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"加载配置文件失败: {e}")
            # 返回默认配置
            return {
                "api_keys": {
                    "openai": "your_openai_api_key_here",
                    "anthropic": "your_anthropic_api_key_here"
                },
                "story_settings": {
                    "max_words": 1000,
                    "age_group": "4-8",
                    "language": "zh-CN"
                },
                "image_settings": {
                    "style": "cartoon",
                    "resolution": "1024x1024",
                    "number_of_images": 4
                }
            }
            
    def generate_images(self, image_prompts):
        """根据提示词生成图片（只生成第一条提示词的图片）"""
        print("\n正在生成插画图片...")
        
        try:
            # 只使用第一条提示词生成一张图片
            if not image_prompts or not image_prompts[0]:
                print("没有找到有效提示词，无法生成图片")
                return None
                
            image_results = []
            prompt = image_prompts[0]  # 只使用第一条提示词
            print(f"正在生成图片...")
            image_path = self.image_generator.generate_image(prompt, download=True, filename="story_image.png")
            image_results.append((prompt, image_path))
            print(f"\n已成功生成 1 张图片")
            return image_results
            
        except Exception as e:
            print(f"生成图片时出错: {str(e)}")
            print("提示：请确保已正确配置API密钥")
            return None
    
    def generate_storybook(self, story_params=None):
        """生成完整的故事绘本（只生成一条提示词和一张图片）
        
        参数:
        - story_params: 故事参数，包含主题、主角、背景等信息。如果为None，则通过命令行获取
        
        返回:
        - dict: 包含故事内容、插画提示词和图片路径的字典
        """
        # 初始化结果字典
        result = {
            "success": False,
            "story_content": "",
            "image_prompts": [],
            "image_exists": False
        }
        
        try:
            # 获取故事参数
            if story_params is None:
                # 兼容命令行模式，通过命令行获取用户输入
                print("=== 欢迎使用儿童故事绘本生成器 ===")
                print("请输入以下信息来生成您的故事绘本：")
                
                theme = input("故事主题（例如：友谊、勇敢、环保等）: ")
                main_character = input("主角（例如：小兔子、小熊、小猫咪等）: ")
                setting = input("故事背景（例如：森林、海洋、太空等）: ")
                moral = input("故事寓意（可选）: ")
                
                story_params = {
                    "theme": theme,
                    "main_character": main_character,
                    "setting": setting,
                    "moral": moral
                }
            else:
                # 从前端接收参数时的日志
                print("\n从前端接收故事参数...")
                print(f"主题: {story_params.get('theme', '未提供')}")
                print(f"主角: {story_params.get('main_character', '未提供')}")
                print(f"背景: {story_params.get('setting', '未提供')}")
                print(f"寓意: {story_params.get('moral', '无')}")
            
            # 验证必要的参数
            required_params = ['theme', 'main_character', 'setting']
            for param in required_params:
                if not story_params.get(param):
                    print(f"错误：缺少必要参数 '{param}'")
                    result["message"] = f"缺少必要参数 '{param}'"
                    return result
            
            print("\n正在生成故事内容...")
            # 生成故事
            story_content = self.story_generator.generate_story(story_params)
            
            if not story_content:
                print("生成故事内容失败")
                result["message"] = "生成故事内容失败"
                return result
            
            print("\n正在生成插画提示词...")
            # 生成插画提示词
            image_prompts = self.image_prompt_generator.generate_prompts(story_content)
            
            if not image_prompts:
                print("生成插画提示词失败")
                result["message"] = "生成插画提示词失败"
                # 即使没有提示词，也要保存故事内容
                self.save_results(story_content, [])
                result["story_content"] = story_content
                return result
            
            # 只保留第一条提示词
            if len(image_prompts) > 1:
                image_prompts = [image_prompts[0]]
                print("已将提示词限制为只保留第一条")
            
            # 保存结果
            self.save_results(story_content, image_prompts)
            
            # 生成图片
            image_results = self.generate_images(image_prompts)
            
            # 更新返回结果
            result = {
                "success": True,
                "story_content": story_content,
                "image_prompts": image_prompts,
                "image_exists": image_results is not None
            }
            
            # 如果成功生成了图片，添加图片路径
            if image_results:
                result["image_path"] = image_results[0][1]
            
            if story_params is None:  # 命令行模式下显示结果
                print("\n=== 故事绘本生成完成！===\n")
                print("故事内容：")
                print(story_content)
                print("\n插画提示词：")
                for i, prompt in enumerate(image_prompts, 1):
                    print(f"{i}. {prompt}")
                
                # 如果成功生成了图片，显示图片信息
                if image_results:
                    print("\n生成的图片：")
                    for i, (_, image_path) in enumerate(image_results, 1):
                        print(f"{i}. {image_path}")
            
            result["message"] = "故事绘本生成成功"
            return result
        except Exception as e:
            print(f"生成故事绘本时出错: {str(e)}")
            result["message"] = f"生成故事绘本时出错: {str(e)}"
            return result
    
    def save_results(self, story_content, image_prompts):
        """保存生成的结果到文件"""
        try:
            # 使用工具函数保存故事内容
            story_file = save_to_file(story_content, "story.txt", "output")
            if story_file:
                print(f"故事内容已保存到: {story_file}")
            
            # 使用工具函数保存插画提示词
            prompts_text = "\n".join([f"{i}. {prompt}" for i, prompt in enumerate(image_prompts, 1)])
            prompts_file = save_to_file(prompts_text, "image_prompts.txt", "output")
            if prompts_file:
                print(f"插画提示词已保存到: {prompts_file}")
        except Exception as e:
            print(f"保存结果文件时出错: {str(e)}")
    
    def run(self):
        """运行应用程序（命令行模式）"""
        try:
            self.generate_storybook()
        except KeyboardInterrupt:
            print("\n程序已被用户中断")
        except Exception as e:
            print(f"程序运行出错: {e}")

if __name__ == "__main__":
    app = StorybookApp()
    app.run()